from math import sin,cos,pi,sqrt
from numpy import array,cross,dot
from enthought.mayavi import mlab
from v3dsolver.diffusion_manager import DiffusionManager
from v3dsolver.fluid_elements.fluid_elems import Saffman
from v3dsolver.merging_manager import MergingManager
import os

def vortexring_diffusion():
    """
    Tests diffusion by simulating the diffusion of a vortex ring which 
    thickens due to diffusion.
    """
    n = 50 # no. of particles in the ring
    t = 0.0
    tf = 1.0
    dt = 0.01
    elems = []
    for i in range(n):
        x = cos(2*pi*i/n)
        y = sin(2*pi*i/n)
        z = 0.
        vor = array([0.,0.,-1.])#cross(array([x,y,z]),array([0.,0.,-1.]))
        elems += [Saffman(pos=array([x,y,z]),vor=vor,
                          vel=array([0.0,0.0,0.0]),rad=None)]
    gamma0 = 1.0
    dm = DiffusionManager(elems,nu=0.04,dt=dt)
    mm = MergingManager(elems,nu=0.04,dt=dt)
    while t<tf:
        x,y,z,g = [],[],[],[]
        for elem in elems:
            x += [elem.pos[0]]
            y += [elem.pos[1]]
            z += [elem.pos[2]]
            g += [sqrt(dot(elem.vor,elem.vor))]
        totalvor = sum([elem.vor for elem in elems])
        print 'Total vor:'+str(totalvor)
        gamma_avg = sum(g)/float(len(elems))
        print 'Gamma_avg:'+str(gamma_avg)
        mlab.points3d(x,y,z,g,scale_mode='none',scale_factor=0.05,
                      vmax=gamma0,colormap='jet')
        filename = os.path.join(os.path.abspath('results'),
                                't = '+str(t)+' sec.png')
        mlab.view(90,90)
        mlab.roll(90)
        mlab.savefig(filename)
        
        dm.diffuse()
        print '\nt = '+str(t)+' sec'
        print 'n_elems after diffusion:'+str(len(elems))
        mm.merge()
        mm.eat()
        print 'n_elems after merging and eating:'+str(len(elems))

        t += dt

if __name__=="__main__":
    vortexring_diffusion()
